'''
该 py 用于将 llm 生成的 markdown(txt) 概要文档 转换为 docx
'''
import os
import re
from docx import Document
import docx
import json
import markdown

from docx.shared import Pt
from docx.enum.table import WD_CELL_VERTICAL_ALIGNMENT, WD_ALIGN_VERTICAL
import argparse


def get_header_level(line):
    # 使用正则表达式匹配不同级别的标题
    # pattern = r'^(\d+(\.\d+)*)'
    pattern = r'\d+(\.\d+)'
    match = re.match(pattern, line)
    if match:
        header_text = match.group(1)
        header_level = line.count('.') + 1
        # print("")
        # print(line)
        # print("### ", header_level, " ###——————", line)
        # print("")
        return header_level
    # elif line.startswith("第"):
    elif line.startswith("#"):
        return 1  # 如果是类似于"第 x 章"的章节标题，设为第 1 级标题
    # print("*** ", " ***——————", line)
    return -1


def make_rows_bold(row):
    for cell in row.cells:
        for paragraph in cell.paragraphs:
            for run in paragraph.runs:
                run.font.bold = True

def make_rows_align_center(row):
    for cell in row.cells:
        paragraphs = cell.paragraphs
        for p in paragraphs:
            p.alignment = docx.enum.text.WD_ALIGN_PARAGRAPH.CENTER

def markdown_table_to_docx(doc, md_table):
    # Parse the Markdown table to HTML with tables extension
    html_table = markdown.markdown(md_table, output_format='html', extensions=['tables'])

    # Split HTML table into header and body sections
    header, body = html_table.split('<tbody>')
    header_cells = header.split('<th>')[1:]

    # Determine number of columns
    cols = len(header_cells)

    # Add a table to the document
    table = doc.add_table(rows=1, cols=cols,style= 'Table Grid')

    # Style the table
    for row in table.rows:
        for cell in row.cells:
            cell.vertical_alignment = WD_CELL_VERTICAL_ALIGNMENT.CENTER
            cell.paragraphs[0].paragraph_format.alignment = WD_ALIGN_VERTICAL.CENTER
            for paragraph in cell.paragraphs:
                for run in paragraph.runs:
                    run.font.size = Pt(10)

    # Parse header cells and add to the table
    header_cells = header.split('<th>')[1:]
    for col_num, cell_text in enumerate(header_cells):
        cell = table.cell(0, col_num)
        cell.text = cell_text.split('</th>')[0].strip()
        cell.paragraphs[0].runs[0].bold = True
        cell.paragraphs[0].paragraph_format.alignment = WD_ALIGN_VERTICAL.CENTER
        cell.vertical_alignment = WD_CELL_VERTICAL_ALIGNMENT.CENTER

    # Parse body cells and add to the table
    for row_num, row in enumerate(body.split('<tr>')[1:]):
        cells = row.split('<td>')[1:]
        row_cells = table.add_row().cells
        for col_num, cell in enumerate(cells):
            cell_text = cell.split('</td>')[0].strip()
            row_cells[col_num].text = cell_text

def handle_markdown_output(input_path):
    with open(input_path, 'r', encoding='utf-8') as f:
        content_list = f.readlines()

    print(content_list)

    for idx, n in enumerate(content_list):
        if n.startswith("#"):
            print(n)
            if check_has_hash_and_space(n):
                # print(f'字符串 "{n}" 以 # 后跟空格开头')
                pass
            else:
                # print(f'字符串 "{n}" 不以 # 后跟空格开头')
                content_list[idx] = insert_space_after_hash(n)
                # content_list[idx] = "\n" + insert_space_after_hash(n)
                if idx != 0 and content_list[idx - 1] not in ['\n', '']:
                    print("content_list[idx-1]:", content_list[idx - 1])
                    content_list[idx] = "\n" + content_list[idx]

        elif n.startswith("表：") or n.startswith("表："):
            print(f'n, ', n)
            if idx != 0 and content_list[idx - 1] not in ['\n', '']:
                print("content_list[idx-1]:", content_list[idx - 1])
                content_list[idx] = "\n" + content_list[idx]

    name = content_list[0].split("#")[-1].strip()
    output_str = ''.join(content_list)

    # with open(name + '_.txt', 'w', encoding='utf-8') as f:
    #     f.write(output_str)

    with open(input_path, 'w', encoding='utf-8') as f:
        f.write(output_str)

pattern = re.compile(r'^#+\s+')
def check_has_hash_and_space(s):
    return bool(pattern.match(s))


pattern2 = re.compile(r'^(#+)(\S)')
def insert_space_after_hash(s):
    return pattern2.sub(r'\1 \2', s)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_files', default='生产制造分册_SOP_生成设计.txt', type=str)
    parser.add_argument('--input_prefix', default='../md/', type=str)
    parser.add_argument('--output_prefix', default='../docx/', type=str)
    args = parser.parse_args()
    
    path = os.path.join(args.input_prefix, args.input_files)
    out_path = args.output_prefix
    
    # 先做表格的预处理
    handle_markdown_output(path)
    
    # for ix, pppp in enumerate(pth):
    # 读取 txt 文件内容
    with open(path, "r", encoding="utf-8") as file:
        lines = file.readlines()

    nam = os.path.basename(path).split(".txt")[0]
    # 创建一个新的 docx 文档
    doc = Document()

    # special_line_started = False
    # is_table_before_this_line = False
    # special_line_content = []

    table_flag = 0
    next_line_flag = 0
    table_list = []
    print("lines:", lines)
    for idddx, line in enumerate(lines):
        line = line.strip()
        # print("line", )
        if idddx == 0:
            line_ = line.split("# ")[-1]
            # doc.add_heading(line_, level=0)
            doc.add_paragraph(line_, style=f'Heading {1}')
            continue

        # 如果是表格
        if line.startswith("表：") or line.startswith("表：") or line.startswith("表：") or line.startswith("表： "):
            table_list = []
            table_flag = 1
            doc.add_paragraph(line)
            continue
        elif line.startswith("|") and table_flag == 1 and next_line_flag == 0:
            table_list.append(line)
            if not line.strip().endswith("|"):
                next_line_flag = 1
            continue
        elif (not line.startswith("|")) and table_flag == 1 and next_line_flag == 0:
            table_flag = 0
            table_str = "\n".join(table_list)
            print("table_str", table_str)
            markdown_table_to_docx(doc, table_str)
            doc.add_paragraph("")
            continue
        elif not line.startswith("|") and table_flag == 1 and next_line_flag == 1:
            print("line: ", line)
            table_list[-1] = table_list[-1].strip() + line

            if line.strip().endswith("|"):
                next_line_flag = 0
            continue
        else:
            header_level = get_header_level(line)

            # 创建相应级别的标题
            if header_level > 0:
                # doc.add_heading(line, level=header_level)
                line_ = line.split("# ")[-1]
                doc.add_paragraph(line_, style=f'Heading {header_level+1}')
            elif line.startswith("[START]") or line.startswith("图")or line.startswith("\n") or line.startswith("\\n") or line.startswith(""):
                continue
            else:
                # 正文内容
                doc.add_paragraph(line)

    # 保存为 docx 文件
    doc.save(os.path.join(out_path, nam + ".docx"))
